# -*- coding: utf-8 -*-
from collections import namedtuple

Token = namedtuple('Token', ['type', 'value'])

class ASTNode:
    def __init__(self, nodetype, value=None, children=None):
        self.type = nodetype
        self.value = value
        self.children = children or []
    def __repr__(self, level=0):
        indent = '  ' * level
        s = f"{indent}{self.type}"
        if self.value is not None:
            s += f": {self.value}"
        s += '\n'
        for c in self.children:
            s += c.__repr__(level+1)
        return s

class ParserError(Exception):
    pass

class RDParser:
    def __init__(self, tokens):
        self.tokens = tokens
        self.pos = 0

    @property
    def cur(self):
        return self.tokens[self.pos]

    def peek(self, n=1):
        idx = self.pos + n
        return self.tokens[idx] if idx < len(self.tokens) else Token('EOF','')

    def eat(self, *, ttype=None, val=None):
        tok = self.cur
        if val is not None:
            if tok.value != val:
                raise ParserError(f"Expected '{val}', got '{tok.value}'")
        elif ttype is not None:
            if tok.type != ttype:
                raise ParserError(f"Expected type {ttype}, got {tok.type}")
        self.pos += 1
        return tok

    def parse(self):
        root = self.parse_Program()
        self.eat(val='')
        return root

    def parse_Program(self):
        decls = []
        while self.cur.type=='KEYWORD' and self.cur.value in ('int','float','char','void'):
            decls.append(self.parse_Decl())
        return ASTNode('Program', children=decls)

    def parse_Decl(self):
        if (self.cur.type=='KEYWORD' and self.peek().type=='ID'
            and self.peek(2).value=='('):
            return self.parse_FunDecl()
        else:
            return self.parse_VarDecl()

    def parse_VarDecl(self):
        # VarDecl → Type InitDecl {',' InitDecl} ';'
        type_tok = self.eat(ttype='KEYWORD')
        children = [ASTNode('Type', type_tok.value)]
        children.append(self.parse_InitDecl())
        while self.cur.value==',':
            self.eat(val=',')
            children.append(self.parse_InitDecl())
        self.eat(val=';')
        return ASTNode('VarDecl', children=children)

    def parse_InitDecl(self):
        idn = self.eat(ttype='ID')
        node = ASTNode('ID', idn.value)
        if self.cur.value=='=':
            self.eat(val='=')
            expr = self.parse_Expr()
            return ASTNode('InitDecl', children=[node, expr])
        return ASTNode('InitDecl', children=[node])

    def parse_FunDecl(self):
        type_tok = self.eat(ttype='KEYWORD')
        id_tok   = self.eat(ttype='ID')
        self.eat(val='(')
        params = self.parse_Params()
        self.eat(val=')')
        body = self.parse_CompoundStmt()
        return ASTNode('FunDecl', children=[
            ASTNode('Type', type_tok.value),
            ASTNode('ID', id_tok.value),
            params, body
        ])

    def parse_Params(self):
        if self.cur.value==')':
            return ASTNode('Params')
        if self.cur.value=='void':
            tok = self.eat(ttype='KEYWORD')
            return ASTNode('Params', value=tok.value)
        nodes = [self.parse_Param()]
        while self.cur.value==',':
            self.eat(val=',')
            nodes.append(self.parse_Param())
        return ASTNode('ParamList', children=nodes)

    def parse_Param(self):
        t = self.eat(ttype='KEYWORD')
        i = self.eat(ttype='ID')
        return ASTNode('Param', children=[ASTNode('Type', t.value), ASTNode('ID', i.value)])

    def parse_CompoundStmt(self):
        self.eat(val='{')
        decls = []
        while self.cur.type=='KEYWORD' and self.cur.value in ('int','float','char','void'):
            decls.append(self.parse_VarDecl())
        stmts = []
        while not (self.cur.value=='}'):
            stmts.append(self.parse_Stmt())
        self.eat(val='}')
        return ASTNode('CompoundStmt', children=[ASTNode('LocalDecls', children=decls),
                                                 ASTNode('StmtList', children=stmts)])

    def parse_Stmt(self):
        t, v = self.cur.type, self.cur.value
        # —— 新增：类型关键字也当作声明语句 —— 
        if t=='KEYWORD' and v in ('int','float','char','void'):
            return self.parse_VarDecl()

        if t=='KEYWORD':
            if v=='return':   return self.parse_ReturnStmt()
            if v=='if':       return self.parse_IfStmt()
            if v=='while':    return self.parse_WhileStmt()
            if v=='do':       return self.parse_DoWhileStmt()
            if v=='for':      return self.parse_ForStmt()
            if v=='break':    return self.parse_BreakStmt()
            if v=='continue': return self.parse_ContinueStmt()
            if v=='switch':   return self.parse_SwitchStmt()
        if t=='SYM' and v=='{':
            return self.parse_CompoundStmt()
        return self.parse_ExprStmt()
    
    def parse_ReturnStmt(self):
        self.eat(val='return')
        if self.cur.value==';':
            self.eat(val=';')
            return ASTNode('ReturnStmt')
        expr = self.parse_Expr()
        self.eat(val=';')
        return ASTNode('ReturnStmt', children=[expr])
    
    def parse_ExprStmt(self):
        if self.cur.value==';':
            self.eat(val=';')
            return ASTNode('ExprStmt')
        e = self.parse_Expr()
        self.eat(val=';')
        return ASTNode('ExprStmt', children=[e])

    def parse_IfStmt(self):
        self.eat(val='if'); self.eat(val='(')
        cond = self.parse_Expr(); self.eat(val=')')
        then = self.parse_Stmt()
        els = None
        if self.cur.value=='else':
            self.eat(val='else')
            els = self.parse_Stmt()
        c = [cond, then]
        if els: c.append(els)
        return ASTNode('IfStmt', children=c)

    def parse_WhileStmt(self):
        self.eat(val='while'); self.eat(val='(')
        cond = self.parse_Expr(); self.eat(val=')')
        body = self.parse_Stmt()
        return ASTNode('WhileStmt', children=[cond, body])

    def parse_DoWhileStmt(self):
        self.eat(val='do')
        body = self.parse_Stmt()
        self.eat(val='while'); self.eat(val='(')
        cond = self.parse_Expr(); self.eat(val=')'); self.eat(val=';')
        return ASTNode('DoWhileStmt', children=[body, cond])

    def parse_ForStmt(self):
        self.eat(val='for')
        self.eat(val='(')
        # ---- init 部分：可能是声明，也可能是表达式 ----
        if self.cur.value == ';':
            init = None
        elif self.cur.type == 'KEYWORD' and self.peek().type == 'ID':
            # 形如: int i = ...;
            init = self.parse_VarDecl()
        else:
            init = self.parse_Expr()
            self.eat(val=';')
         # 处理 condition 和 step 如前
        cond = None
        if self.cur.value!=';':
            cond = self.parse_Expr()
        self.eat(val=';')
        step = None
        if self.cur.value!=')':
            step = self.parse_Expr()
        self.eat(val=')')
        body = self.parse_Stmt()
        return ASTNode('ForStmt', children=[init, cond, step, body])

    def parse_BreakStmt(self):
        self.eat(val='break'); self.eat(val=';')
        return ASTNode('BreakStmt')

    def parse_ContinueStmt(self):
        self.eat(val='continue'); self.eat(val=';')
        return ASTNode('ContinueStmt')

    def parse_SwitchStmt(self):
        self.eat(val='switch'); self.eat(val='(')
        expr = self.parse_Expr(); self.eat(val=')')
        self.eat(val='{')
        cases = []
        default = None
        while self.cur.value=='case':
            cases.append(self.parse_CaseClause())
        if self.cur.value=='default':
            default = self.parse_DefaultClause()
        self.eat(val='}')
        children = [expr] + cases
        if default: children.append(default)
        return ASTNode('SwitchStmt', children=children)

    def parse_CaseClause(self):
        self.eat(val='case')
        val = self.parse_Expr(); self.eat(val=':')
        stmts = []
        while self.cur.value not in ('case','default','}'):
            stmts.append(self.parse_Stmt())
        return ASTNode('CaseClause', children=[val] + stmts)

    def parse_DefaultClause(self):
        self.eat(val='default'); self.eat(val=':')
        stmts = []
        while self.cur.value not in ('case','}','default'):
            stmts.append(self.parse_Stmt())
        return ASTNode('DefaultClause', children=stmts)

    def parse_Expr(self):
        if self.cur.type=='ID' and self.peek().value=='=':
            idn = ASTNode('ID', self.eat(ttype='ID').value)
            self.eat(val='=')
            rhs = self.parse_Expr()
            return ASTNode('Assign', children=[idn, rhs])
        return self.parse_LogicalOr()

    def parse_LogicalOr(self):
        node = self.parse_LogicalAnd()
        while self.cur.value=='||':
            op = ASTNode('LogicalOr', self.eat(val='||').value)
            rhs = self.parse_LogicalAnd()
            node = ASTNode('BinaryOp', children=[node, op, rhs])
        return node

    def parse_LogicalAnd(self):
        node = self.parse_RelExpr()
        while self.cur.value=='&&':
            op = ASTNode('LogicalAnd', self.eat(val='&&').value)
            rhs = self.parse_RelExpr()
            node = ASTNode('BinaryOp', children=[node, op, rhs])
        return node

    def parse_RelExpr(self):
        left = self.parse_AddExpr()
        if self.cur.value in ('<','<=','>','>=','==','!='):
            op = ASTNode('RelOp', self.eat(val=self.cur.value).value)
            right = self.parse_AddExpr()
            return ASTNode('BinaryOp', children=[left, op, right])
        return left

    def parse_AddExpr(self):
        node = self.parse_Term()
        while self.cur.value in ('+','-'):
            op = ASTNode('AddOp', self.eat(val=self.cur.value).value)
            rhs = self.parse_Term()
            node = ASTNode('BinaryOp', children=[node, op, rhs])
        return node

    def parse_Term(self):
        node = self.parse_Factor()
        while self.cur.value in ('*','/','%'):
            op = ASTNode('MulOp', self.eat(val=self.cur.value).value)
            rhs = self.parse_Factor()
            node = ASTNode('BinaryOp', children=[node, op, rhs])
        return node

    def parse_Factor(self):
        if self.cur.value in ('+','-'):
            op = self.eat(val=self.cur.value)
            return ASTNode('UnaryOp', value=op.value, children=[self.parse_Factor()])
        if self.cur.value=='(':
            self.eat(val='(')
            n = self.parse_Expr()
            self.eat(val=')')
            return n
        if self.cur.type=='ID' and self.peek().value=='(':
            name = self.eat(ttype='ID').value
            self.eat(val='(')
            args = self.parse_ArgumentList()
            self.eat(val=')')
            return ASTNode('Call', value=name, children=[args])
        if self.cur.type=='STRING':
            return ASTNode('STRING', self.eat(ttype='STRING').value)
        if self.cur.type=='ID':
            return ASTNode('ID', self.eat(ttype='ID').value)
        if self.cur.type=='NUM':
            return ASTNode('NUM', self.eat(ttype='NUM').value)
        raise ParserError(f"Unexpected token in Factor: {self.cur.value}")

    def parse_ArgumentList(self):
        args = []
        if self.cur.value==')':
            return ASTNode('ArgList', children=[])
        args.append(self.parse_Expr())
        while self.cur.value==',':
            self.eat(val=',')
            args.append(self.parse_Expr())
        return ASTNode('ArgList', children=args)

def run_parser(tokens):
    tokens.append(Token('EOF',''))
    return RDParser(tokens).parse()
